package com.example.sharding.jdbc.demo.config;

import com.example.sharding.jdbc.demo.config.properties.ShardingJdbcDatasourceProperties;
import com.example.sharding.jdbc.demo.config.properties.ShardingJdbcRuleProperties;
import com.example.sharding.jdbc.demo.config.properties.YamlShardingDatasourceProperties;
import com.example.sharding.jdbc.demo.constants.DSConstants;
import com.example.sharding.jdbc.demo.constants.ShardingConstants;
import com.zaxxer.hikari.HikariDataSource;
import org.apache.shardingsphere.driver.api.ShardingSphereDataSourceFactory;
import org.apache.shardingsphere.infra.config.RuleConfiguration;
import org.apache.shardingsphere.infra.config.algorithm.ShardingSphereAlgorithmConfiguration;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.strategy.sharding.StandardShardingStrategyConfiguration;
import org.apache.shardingsphere.sharding.yaml.config.rule.YamlTableRuleConfiguration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;

/**
 * ShardingJdbc配置
 *
 * @author airtrioa
 * @date 2021/4/30
 */
@Configuration
public class ShardingJdbcConfig {
  private static final Logger logger = LoggerFactory.getLogger(ShardingJdbcConfig.class);
  private final Map<String, DataSource> dataSourceMap;
  private final Map<String, RuleConfiguration> ruleMap;

  public ShardingJdbcConfig(ShardingJdbcRuleProperties shardingJdbcProperties,
                            ShardingJdbcDatasourceProperties shardingJdbcDatasourceProperties) {
    this.dataSourceMap = this.getShardingDatasourceMap(shardingJdbcDatasourceProperties);
    this.ruleMap = this.getShardingTableRuleMap(shardingJdbcProperties);
  }

  /**
   * 创建分片数据源
   *
   * @return 数据源
   * @throws SQLException SQL 异常
   */
  @Bean(name = DSConstants.TEST_LOG_DS_NAME)
  public DataSource billLogSharding() throws SQLException {
    return ShardingSphereDataSourceFactory.createDataSource(this.dataSourceMap,
        Collections.singleton(ruleMap.get(DSConstants.TEST_LOG_DS_NAME)),
        new Properties());
  }

  /**
   * 创建分片数据源
   *
   * @return 数据源
   * @throws SQLException SQL 异常
   */
  @Bean(name = DSConstants.TEST_RECORD_DS_NAME)
  public DataSource recordSharding() throws SQLException {
    return ShardingSphereDataSourceFactory.createDataSource(this.dataSourceMap,
        Collections.singleton(ruleMap.get(DSConstants.TEST_RECORD_DS_NAME)),
        new Properties());
  }

  /**
   * 获取数据源Map
   *
   * @param shardingJdbcDatasourceProperties 数据源参数
   * @return 分片数据源Map
   */
  private Map<String, DataSource> getShardingDatasourceMap(ShardingJdbcDatasourceProperties shardingJdbcDatasourceProperties) {
    // 获取数据配置 并构建 分片数据源
    Map<String, YamlShardingDatasourceProperties> datasourceMap = shardingJdbcDatasourceProperties.getDatasource();
    logger.info("开始构建分片数据源列表:[{}]", datasourceMap.keySet());
    // 配置真实数据源
    Map<String, DataSource> stringDataSourceHashMap = new HashMap<>(datasourceMap.size());
    datasourceMap.forEach((dsName, dsConfig) -> {
      HikariDataSource ivrDataSource = new HikariDataSource();
      ivrDataSource.setDriverClassName(dsConfig.getDriverClassName());
      ivrDataSource.setJdbcUrl(dsConfig.getJdbcUrl());
      ivrDataSource.setUsername(dsConfig.getUn());
      ivrDataSource.setPassword(dsConfig.getPwd());
      ivrDataSource.setMaximumPoolSize(dsConfig.getMaxPoolSize());
      stringDataSourceHashMap.put(dsName, ivrDataSource);
    });
    return stringDataSourceHashMap;
  }

  /**
   * 获取分片配置映射
   *
   * @param shardingJdbcRuleProperties 分片规则参数
   * @return 分片配置映射
   */
  private Map<String, RuleConfiguration> getShardingTableRuleMap(ShardingJdbcRuleProperties shardingJdbcRuleProperties) {
    Map<String, YamlTableRuleConfiguration> tables = shardingJdbcRuleProperties.getTables();
    logger.info("需要进行分片的表集合:[{}]", tables.keySet());
    Map<String, RuleConfiguration> configurationHashMap = new HashMap<>(tables.size());
    tables.forEach((logicTable, tableConfig) -> {
      logger.info("开始构建表:[{}]分片算法", logicTable);
      // 配置 表规则
      ShardingTableRuleConfiguration tableRuleConfig = new ShardingTableRuleConfiguration(logicTable, tableConfig.getActualDataNodes());
      // 配置分表策略
      // 设置分片为标准分片且指定分片键
      String shardingColumn = tableConfig.getTableStrategy().getStandard().getShardingColumn();
      tableRuleConfig.setTableShardingStrategy(new StandardShardingStrategyConfiguration(shardingColumn, ShardingConstants.TABLE_SHARDING_ALGORITHM_KEY));

      // 配置分片规则
      ShardingRuleConfiguration shardingRuleConfig = new ShardingRuleConfiguration();
      shardingRuleConfig.getTables().add(tableRuleConfig);

      // 配置分表算法 设置按月分表
      shardingRuleConfig.getShardingAlgorithms().put(ShardingConstants.TABLE_SHARDING_ALGORITHM_KEY,
          new ShardingSphereAlgorithmConfiguration(ShardingConstants.MONTH_SHARD, new Properties()));
      configurationHashMap.put(logicTable, shardingRuleConfig);
    });
    return configurationHashMap;
  }
}
